# -*- coding: utf-8 -*-
# @File:     proccess_generate_result.py
# @Author:
# @DateTime: 2025/06/10/16:45
import asyncio
import traceback

from app.api.generate.utils.redis_utils import update_task_state
from app.db.session import SessionLocal
from app.models.generate import GenerateTask, GenerateRecord, GenerateResult, UserMembership
from app.utils.file_storage import get_storage_file
from app.utils.logger import logger


def process_generate_result(message_data):
    print(message_data)
    task_id = message_data.get('task_id')
    result_list = message_data.get('result', [])
    result_type = message_data.get('result_type', 'png')
    resp = message_data.get('resp', {})
    user_id = resp.get('user_id')
    message = message_data.get('message')

    membership_id = resp.get('membership')
    point = resp.get('point', 1)

    db = SessionLocal()

    try:
        task = db.query(GenerateTask).filter(GenerateTask.id == task_id, GenerateTask.state == 2).first()
        record = db.query(GenerateRecord).filter(GenerateRecord.g_task_id == task_id if task else None, GenerateRecord.state == 2).first()
    except:
        task = None
        record = None
    try:
        if result_list:

            try:
                membership = db.query(UserMembership).filter(UserMembership.id == membership_id).first()
                if membership.total - membership.use_number > point:
                    state = 3
                else:
                    state = -1
                    message = '生成点数不足'
                membership.use_number = membership.use_number + point
                db.add(membership)
                db.commit()
                db.refresh(membership)
            except:
                logger.error(f"验证套餐错误：\n{traceback.format_exc()}")
                state = -1
                message = '生成点数不足'

            li = []
            for result in result_list:
                path = result['path']
                file_type = result['type']
                data = {
                    'task_id': task.id if state == 3 else None,
                    'record_id': record.id if record else None,
                    'path': path,
                    'file_type': file_type,
                }
                db_obj = GenerateResult(**data)  # type: ignore
                db.add(db_obj)
                db.commit()
                db.refresh(db_obj)

                path_url = get_storage_file(path)
                li.append({
                    'id': db_obj.id,
                    'path': path,
                    'path_url': path_url,
                    'file_type': file_type,
                    'download': db_obj.download,
                })

            if task:
                asyncio.run(update_task_state(task_id, state, li))
                task.state = state
                task.error_describe = message
                db.add(task)
                db.commit()
                db.refresh(task)
            if record:
                record.state = 3
                db.add(record)
                db.commit()
                db.refresh(record)
        else:
            if task:
                asyncio.run(update_task_state(task_id, -1, []))
                task.error_describe = message
                task.state = -1
                db.add(task)
                db.commit()
                db.refresh(task)
            if record:
                record.error_describe = message
                record.state = -1
                db.add(record)
                db.commit()
                db.refresh(record)
    except:
        if task:
            asyncio.run(update_task_state(task_id, -1, []))
            task.error_describe = message
            task.state = -1
            db.add(task)
            db.commit()
            db.refresh(task)
        if record:
            record.error_describe = message
            record.state = -1
            db.add(record)
            db.commit()
            db.refresh(record)
    finally:
        db.close()

